import logging
import os
from common.config import get_logger
import numpy as np
import tensorflow as tf
from keras.applications.resnet50 import ResNet50
from keras.applications.resnet import ResNet101
from keras.applications.resnet import ResNet152
from keras.applications.xception import Xception
from keras.applications.resnet50 import preprocess_input
from keras.preprocessing import image
from numpy import linalg as LA

os.environ['KERAS_HOME'] = os.path.abspath(os.path.join('.', 'data'))

logging = get_logger("resnet50")


class CustomOperator:
    def __init__(self):
        self.model = ResNet50(weights='imagenet', include_top=False, pooling='avg')

    def execute(self, img_path, ignore_error_img=False):
        try:
            img = image.load_img(img_path, target_size=(224, 224))
            x = image.img_to_array(img)
            x = np.expand_dims(x, axis=0)
            x = preprocess_input(x)
            features = self.model.predict(x)
            norm_feature = features[0] / LA.norm(features[0])
            norm_feature = [i.item() for i in norm_feature]
            # logging.debug(f"ResNet50 execute, img_path: {img_path}, norm_feature size: {len(norm_feature)}")
            return norm_feature
        except Exception as e:
            logging.error(f"ResNet50, {str(img_path)} execute error: {e}")
            if ignore_error_img:
                return None
            else:
                raise RuntimeError(str(img_path)) from e
